import numpy as np
from smt.sampling_methods import Random
from smt.sampling_methods import FullFactorial
from typing import List, Tuple, Optional, Callable, Dict, Sequence
import os


class Sampler(object):
    """
    统一采样器 (Python 3.7 兼容版）
    ----------------------------------------------------
    只做三件事：接收 bounds -> 生成样本 -> （可选）保存文件

    内置方法：
      - 'random'         : 随机采样
      - 'lhs'            : 拉丁超立方采样
      - 'fullfactorial'  : 全因子采样 / 格点采样

    可通过 register_method(name, fn) 注册自定义方法。
    """

    _registry = {}  # type: Dict[str, Callable[['Sampler'], np.ndarray]]

    def __init__(
        self,
        bounds,                         # type: Sequence[Tuple[float, float]]
        method="lhs",                   # type: str
        n_samples=50,                   # type: int
        seed=None,                      # type: Optional[int]
    ):
        # 校验/缓存
        # type: List[Tuple[float, float]]
        self.bounds = np.asarray(self._validate_bounds(bounds), dtype=float)
        # type: int
        self.dim = len(self.bounds)
        self.method = str(method).lower().strip(
        )                                   # type: str
        # type: int
        self.n_samples = int(n_samples)
        # type: Optional[int]
        self.seed = seed

        if self.n_samples <= 0:
            raise ValueError("n_samples 必须是正整数。")

    # ---------------- Public API ----------------

    def sample(self):
        # type: () -> np.ndarray
        """主入口：返回 (N, D) 数组。"""
        m = self.method
        if m in self._registry:
            X = self._registry[m](self)
        elif m == "random":
            X = self._sample_random()
        elif m == "lhs":
            X = self._sample_lhs()
        elif m == "fullfactorial":
            X = self._sample_fullfactorial()
        else:
            raise ValueError(
                "Unknown method '%s'. Supported: random, lhs, fullfactorial, mc (+ registered)." % m
            )
        return X

    def save_txt(self, path, data=None, float_fmt="%.6f"):
        """
        把采样结果保存为纯数字 TXT (无表头、无元信息，每行一条样本，空格分隔）
        """
        X = data if data is not None else self.sample()
        if X.ndim != 2 or X.shape[1] != self.dim:
            raise ValueError("data 必须是 (N, D) 且 D == 维度数。")
        # 确保目录存在
        d = os.path.dirname(os.path.abspath(path)) or "."
        if not os.path.exists(d):
            raise FileNotFoundError("目标目录不存在: %s" % d)

        fmt = " ".join([float_fmt] * self.dim)
        with open(path, "w", encoding="utf-8") as f:
            for row in X:
                f.write((fmt % tuple(row)) + "\n")
        return path

    # ---------------- Method registry ----------------

    @classmethod
    def register_method(cls, name, fn):
        # type: (str, Callable[['Sampler'], np.ndarray]) -> None
        key = str(name).lower().strip()
        if key in {"random", "lhs", "fullfactorial"}:
            raise ValueError("'%s' 是内置方法名，请换一个名字。" % name)
        cls._registry[key] = fn

    # ---------------- Built-in implementations ----------------
    def _sample_random(self):
        sampling = Random(xlimits=self.bounds)
        x = sampling(self.n_samples)
        return x

    def _sample_lhs(self):
        lhd = np.zeros((self.n_samples, self.dim))
        for i in range(self.dim):
            perm = np.random.permutation(self.n_samples) + 1
            lhd[:, i] = (perm - np.random.rand(self.n_samples)) / \
                self.n_samples
        # 将样本缩放到提供的限制范围
        x = np.zeros_like(lhd)
        for i in range(self.dim):
            x[:, i] = self.bounds[i, 0] + lhd[:, i] * \
                (self.bounds[i, 1] - self.bounds[i, 0])
        return x

    def _sample_fullfactorial(self):
        sampling = FullFactorial(xlimits=self.bounds)
        x = sampling(self.n_samples)
        return x

    # ---------------- Utils ----------------

    @staticmethod
    def _validate_bounds(bounds):
        # type: (Sequence[Tuple[float, float]]) -> List[Tuple[float, float]]
        out = []
        for b in bounds:
            if len(b) != 2:
                raise ValueError("每个 bounds 条目必须是 (low, high) 元组。")
            lo, hi = float(b[0]), float(b[1])
            if not (lo < hi):
                raise ValueError("非法边界 %r: 需要 low < high。" % (b,))
            out.append((lo, hi))
        return out


# ----------------- Examples -----------------

# # === 准备边界 (D,2) 或 [(lo,hi), ...] 都行 ===
# bounds = [(-5.0, 5.0), (-2.0, 5.2), (0.0, 1.0)]   # D = 3
# n_samples = 12
# seed = 2025

# # === 输出目录（必须已存在，否则 save_txt 会 FileNotFoundError）===
# out_dir = "D:\\降阶代码重写\\utils"
# os.makedirs(out_dir, exist_ok=True)

# # === 1) Random ===
# sampler_rand = Sampler(bounds, method="random", n_samples=n_samples, seed=seed)
# X_rand = sampler_rand.sample()
# print("Random shape:", X_rand.shape)
# sampler_rand.save_txt(os.path.join(out_dir, "sample_points.txt"), data=X_rand, float_fmt="%.3f")
